package com.leetcode;

import java.util.HashMap;
import java.util.Map;

/**
 * 路径总和
 *
 * @author winson
 * @create 2022-04-24 10:41
 */
public class PathSumSolution {

    public int pathSum(TreeNode root, int targetSum) {
        if (root == null) {
            return 0;
        }
        /**
         * 路径总和 = 当前节点路径 + 当前节点左子节点路径 + 当前节点右子节点路径
         */
        int ans = rootSum(root, targetSum);
        ans += rootSum(root.left, targetSum);
        ans += rootSum(root.right, targetSum);
        return ans;
    }

    public int rootSum(TreeNode root, int targetSum) {
        int ans = 0;
        if (root == null) {
            return ans;
        }
        if (root.val == targetSum) {
            ans++;
        }
        ans += pathSum(root.left, targetSum - root.val);
        ans += pathSum(root.right, targetSum - root.val);
        return ans;
    }

    public int pathSumPrefix(TreeNode root, int targetSum) {
        Map<Long, Integer> prefix = new HashMap<>();
        prefix.put(0L, 1);
        return dfs(root, prefix, 0, targetSum);
    }

    private int dfs(TreeNode root, Map<Long, Integer> prefix, long curr, int targetSum) {
        if (root == null) {
            return 0;
        }
        int ans = 0;
        curr += root.val;
        ans = prefix.getOrDefault(curr - targetSum, 0);
        prefix.put(curr, prefix.getOrDefault(curr, 0) + 1);
        ans += dfs(root.left, prefix, curr, targetSum);
        ans += dfs(root.right, prefix, curr, targetSum);
        prefix.put(curr, prefix.getOrDefault(curr, 0) - 1);
        return ans;
    }
}
